package com.gitee.zycra.future.job.common.util;

import lombok.extern.slf4j.Slf4j;

import javax.crypto.Cipher;
import java.io.ByteArrayOutputStream;
import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.security.Key;
import java.security.KeyFactory;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Signature;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.RSAPrivateKeySpec;
import java.security.spec.RSAPublicKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * @author zycra
 */
@Slf4j
public final class RSAUtil extends BCUtil {

    private static final String RSA_LABEL = "RSA";
    private static final String RSA_CIPHER = "RSA/ECB/OAEPWITHSHA-256ANDMGF1PADDING";
    private static final String RSA_SIGN_ALGORITHMS = "SHA1WithRSA";
    private static final int RSA_SHIFT = 3;
    private static final int RSA_IN_BLOCK_SUBTRACT = 11;

    private static final String PUBLIC_KEY_HEAD = "-----BEGIN PUBLIC KEY-----";

    private static final String PUBLIC_KEY_TAIL = "-----END PUBLIC KEY-----";

    private static final String PRIVATE_KEY_HEAD = "-----BEGIN PRIVATE KEY-----";

    private static final String PRIVATE_KEY_TAIL = "-----END PRIVATE KEY-----";

    private static final String NEXT_LINE = "\n";

    private static final int KEY_BLOCK_LENGTH = 64;

    private RSAUtil() {
    }

    /**
     * RSA公钥加密
     *
     * @param content   待加密的内容
     * @param publicKey 公钥字符串
     * @return 加密后的内容
     */
    public static String encryptByPublicKey(String content, String publicKey) {
        return doRSA(content, publicKey, true);
    }

    /**
     * RSA私钥解密
     *
     * @param content    待解密的内容
     * @param privateKey 私钥字符串
     * @return 解密后的内容
     */
    public static String decryptByPrivateKey(String content, String privateKey) {
        return doRSA(content, privateKey, false);
    }

    private static PrivateKey getPrivateKey(String key) {
        if (key == null) {
            return null;
        }
        try {
            byte[] keyBytes = Base64.getMimeDecoder().decode(key);
            return KeyFactory.getInstance(RSA_LABEL).generatePrivate(new PKCS8EncodedKeySpec(keyBytes));
        } catch (Exception e) {
            log.error("get private key error", e);
        }
        return null;
    }

    private static PublicKey getPublicKey(String key) {
        if (key == null) {
            return null;
        }
        try {
            byte[] keyBytes = Base64.getMimeDecoder().decode(key);
            return KeyFactory.getInstance(RSA_LABEL).generatePublic(new X509EncodedKeySpec(keyBytes));
        } catch (Exception e) {
            log.error("get public key error", e);
        }
        return null;
    }

    private static int getKeySize(Key key, boolean forEncryption) {
        try {
            KeyFactory keyFactory = KeyFactory.getInstance(RSA_LABEL);
            BigInteger modules = forEncryption ? keyFactory.getKeySpec(key, RSAPublicKeySpec.class).getModulus()
                    : keyFactory.getKeySpec(key, RSAPrivateKeySpec.class).getModulus();
            return modules.bitLength();
        } catch (Exception e) {
            log.error("get key size error", e);
        }
        return -1;
    }

    private static String doRSA(String content, String key, boolean forEncryption) {
        if (content == null) {
            return null;
        }
        Key rsaKey = forEncryption ? getPublicKey(key) : getPrivateKey(key);
        if (rsaKey == null) {
            return null;
        }
        int keySize = getKeySize(rsaKey, forEncryption);
        if (keySize < 0) {
            return null;
        }
        int cipherMode = forEncryption ? Cipher.ENCRYPT_MODE : Cipher.DECRYPT_MODE;
        int maxBlock = forEncryption ? (keySize >> RSA_SHIFT) - RSA_IN_BLOCK_SUBTRACT : keySize >> RSA_SHIFT;
        byte[] data = forEncryption ? content.getBytes(StandardCharsets.UTF_8) : Base64.getMimeDecoder().decode(content);
        int inputLen = data.length;
        int offSet = 0;
        byte[] cache;
        int i = 0;
        try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
            Cipher cipher = Cipher.getInstance(RSA_CIPHER);
            cipher.init(cipherMode, rsaKey);
            while (inputLen - offSet > 0) {
                cache = cipher.doFinal(data, offSet, Math.min(inputLen - offSet, maxBlock));
                out.write(cache, 0, cache.length);
                i++;
                offSet = i * maxBlock;
            }
            byte[] result = out.toByteArray();
            return forEncryption ? Base64.getEncoder().encodeToString(result) : new String(result, StandardCharsets.UTF_8);
        } catch (Exception e) {
            log.error("do RSA error", e);
        }
        return null;
    }

    /**
     * RSA私钥签名
     *
     * @param content    待签名的内容
     * @param privateKey RSA私钥字符串
     * @return 签名后的字符串
     */
    public static String signByPrivateKey(String content, String privateKey) {
        if (content == null) {
            return null;
        }
        PrivateKey priKey = getPrivateKey(privateKey);
        if (priKey == null) {
            return null;
        }
        try {
            Signature signature = Signature.getInstance(RSA_SIGN_ALGORITHMS);
            signature.initSign(priKey);
            signature.update(content.getBytes(StandardCharsets.UTF_8));
            byte[] sign = signature.sign();
            return Base64.getEncoder().encodeToString(sign);
        } catch (Exception e) {
            log.error("do sign by private key error", e);
        }
        return null;
    }

    /**
     * RSA公钥验证签名
     *
     * @param content   待验签的内容
     * @param publicKey RSA公钥字符串
     * @param signStr   待验证的签名字符串
     * @return 验证结果
     */
    public static boolean verifyByPublicKey(String content, String publicKey, String signStr) {
        if (content == null) {
            return false;
        }
        PublicKey pubKey = getPublicKey(publicKey);
        if (pubKey == null) {
            return false;
        }
        try {
            Signature signature = Signature.getInstance(RSA_SIGN_ALGORITHMS);
            signature.initVerify(pubKey);
            signature.update(content.getBytes(StandardCharsets.UTF_8));
            byte[] sign = Base64.getMimeDecoder().decode(signStr);
            return signature.verify(sign);
        } catch (Exception e) {
            log.error("do verify sign by public key error", e);
        }
        return false;
    }

    /**
     * 单行密钥转换为格式化后的密钥
     *
     * @param eyk        单行密钥
     * @param privateKey 是否是私钥
     * @return 格式化后的密钥
     */
    public static String convertToFormatKey(String eyk, boolean privateKey) {
        StringBuilder result = new StringBuilder(privateKey ? PRIVATE_KEY_HEAD : PUBLIC_KEY_HEAD);
        int length = eyk.length();
        int index = 0;
        while (index < length) {
            result.append(NEXT_LINE).append(eyk, index, Math.min(index + KEY_BLOCK_LENGTH, length));
            index += KEY_BLOCK_LENGTH;
        }
        result.append(NEXT_LINE).append(privateKey ? PRIVATE_KEY_TAIL : PUBLIC_KEY_TAIL);
        return result.toString();
    }
}
